# -*- coding: utf-8 -*-
"""
/***************************************************************************
 WaterExtractMTDialog
                                 A QGIS plugin
 Water Extract Multi Thread
                             -------------------
        begin                : 2015-06-23
        copyright            : (C) 2015 by xushiluo
        email                : xushiluo@163.com
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

from PyQt4 import QtCore, QtGui
from PyQt4.QtCore import *

from ui_waterextractmt import Ui_WaterExtractMT
# create the dialog for zoom to point
from qgis.core import QgsMessageLog, QgsRasterLayer, QgsMapLayerRegistry
from qgis.gui import QgsMessageBar
from qgis.utils import iface
from osgeo import gdal
from waterExtractWorker import Worker

class WaterExtractMTDialog(QtGui.QDialog, Ui_WaterExtractMT):
    def __init__(self):
        QtGui.QDialog.__init__(self)
        self.setupUi(self)
        # Set up the signals
        self.connectSignals()

        # Additional code
        self.inFileName = None
        self.outFileName = None
        self.rasterBands = None
        self.thread = None
        self.worker = None

    def connectSignals(self):
        self.cancelButton.clicked.connect(self.close)
        self.inputButton.clicked.connect(self.showOpenDialog)
        self.outputButton.clicked.connect(self.showSaveDialog)
        self.okButton.clicked.connect(self.startWorker)
        self.connect(self.inputComboBox, SIGNAL("currentIndexChanged(int)"), self.setBandsComboBoxByIndex)

    #获取所有的栅格图层
    def getRasterLayers(self):
        layers = QgsMapLayerRegistry.instance().mapLayers().values()
        raster = []
        for layer in layers:
            if layer.type() == layer.RasterLayer:
                if layer.providerType() == 'gdal':  # only gdal file-based layers
                    raster.append(layer)
        return raster

    def getExtendedLayerName(self, layer):
        authid = layer.crs().authid()
        if authid is not None:
            return layer.name() + ' [' + str(authid) + ']'
        else:
            return layer.name()

    # 根据图层对象获取输入图层的路径
    def getInputLayerPath(self, layerObj):
        lyrPath = None
        if isinstance(layerObj, QgsRasterLayer):
            lyrPath = unicode(layerObj.dataProvider().dataSourceUri())
        else:
            lyrPath = unicode(layerObj)
            layers = self.getRasterLayers()
            for layer in layers:
                if layer.name() == lyrPath:
                    lyrPath = unicode(layer.dataProvider().dataSourceUri())
        return lyrPath

    # 根据文件名读取波段数目
    def setBandsComboBox(self, inputFileName):
        gdal.AllRegister()

        dataset = gdal.Open(str(inputFileName))
        self.rasterBands = dataset.RasterCount
        dataset = None

        self.redComboBox.clear()
        self.greenComboBox.clear()
        self.nirComboBox.clear()
        self.mRComboBox.clear()

        for i in range(self.rasterBands, 0, -1):
            self.redComboBox.addItem(str(i))
            self.greenComboBox.addItem(str(i))
            self.nirComboBox.addItem(str(i))
            self.mRComboBox.addItem(str(i))

        self.redComboBox.setDisabled(False)
        self.greenComboBox.setDisabled(False)
        self.nirComboBox.setDisabled(False)
        self.mRComboBox.setDisabled(False)

        #输入栅格文件的路径
        self.inFileName = self.getInputLayerPath( self.inputComboBox.itemData(self.inputComboBox.currentIndex()))

    # 根据组合框中的输入文件索引，读取该栅格文件对应的波段数目
    def setBandsComboBoxByIndex(self, comboBoxIndex):
        lyrObj = self.inputComboBox.itemData( comboBoxIndex )
        if comboBoxIndex!=-1:
            tempFileName = self.getInputLayerPath(lyrObj)
            self.setBandsComboBox(tempFileName)

    # 对话框显示时的事件处理函数
    def showEvent(self, event):
        super(WaterExtractMTDialog, self).showEvent(event)
        rasterLayers = self.getRasterLayers()
        items=[]
        for layer in rasterLayers:
            items.append((self.getExtendedLayerName(layer), layer))
        # 先清空QComboBox的所有对象
        self.inputComboBox.clear()
        for (name, value) in items:
            self.inputComboBox.addItem(name, value)
        if self.inputComboBox.count() >= 1:
            self.inputComboBox.setCurrentIndex(self.inputComboBox.count() - 1)
            lyrObj = self.inputComboBox.itemData(self.inputComboBox.currentIndex())
            tempFileName = self.getInputLayerPath(lyrObj)
            self.setBandsComboBox(tempFileName)

    def showOpenDialog(self):
        fileName = str(QtGui.QFileDialog.getOpenFileName(self, "Input Raster File:"))

        if len(fileName) is 0:
            return
        else:
            self.inFileName = fileName

        if self.inFileName is not None and self.outFileName is not None:
            self.okButton.setDisabled(False)

        self.inputComboBox.clear()
        self.inputComboBox.addItem(self.inFileName,self.inFileName)
        self.inputComboBox.setCurrentIndex(self.inputComboBox.count() - 1)

    def showSaveDialog(self):
        #self.outFileName = str(QtGui.QFileDialog.getSaveFileName(self,
        #        'Output Raster File:', '', '*.tif'))

        # Declare the filetype in which to save the output file
        # Currently the plugin only supports GeoTIFF files
        fileTypes = 'Tiff Files (*.tif *.tiff)'
        fileName, filter = QtGui.QFileDialog.getSaveFileNameAndFilter(
            self, 'Output Raster File:', '', fileTypes)

        if len(fileName) is 0:
            return
        else:
            # Extract the base filename without the suffix if it exists
            # Convert the fileName from QString to python string
            fileNameStr = str(fileName)

            # Split the fileNameStr where/if a '.' exists
            splittedFileName = fileNameStr.split('.')

            # Finally extract the base filename from the splitted filename
            baseFileName = splittedFileName[0]

            # Initialize the suffix string
            suffixStr = ''

            # Check if the user entered a suffix
            suffixExists = False
            existingSuffix = ''
            if len(splittedFileName) != 1:
                existingSuffix = splittedFileName[len(splittedFileName) - 1]
                if existingSuffix is not None:
                    suffixExists = True


            # Extract the suffix from the selected filetype filter
            # Convert the selected filter from QString to python string
            filterStr = str(filter)

            # Split the filter string where/if an asterisk (*) exists
            # I do this to find where the first suffix of the selected filetype
            # occurs
            splittedFilter = filterStr.split('*')

            # If a suffix is not supplied by the user it will be automatically
            # added to the filename. The default suffix will be the first
            # available suffix for the chosen filetype
            if not suffixExists:
                # Extract the 'dirty' suffix string where the first suffix is located
                dirtySuffixStr = splittedFilter[1]

                # Find out the number of the available suffixes
                suffixNum = len(splittedFilter) - 1

                if suffixNum == 1:
                    # Split the dirty suffix string where a ')' occurs
                    # which indicates where the selected filetype ends
                    splittedDirtySuffixStr = dirtySuffixStr.split(')')
                else:
                    # Split the dirty suffix string where a space occurs which
                    # indicates where the selected filetype suffix ends
                    splittedDirtySuffixStr = dirtySuffixStr.split(' ')
                suffixStr = splittedDirtySuffixStr[0]
            else:
                # WE NEED TO CHECK IF THE SUPPLIED SUFFIX CORRESPONDS TO THE
                # SELECTED FILETYPE

                # Extract all the suffixes available for the selected filetype
                # First find out the number of the available suffixes
                suffixNum = len(splittedFilter) - 1

                if suffixNum == 1:
                    # Extract the 'dirty' suffix string where the suffix is located
                    dirtySuffixStr = splittedFilter[1]

                    # Split the dirty suffix string where a space occurs which
                    # indicates where the selected filetype suffix ends
                    splittedDirtySuffixStr = dirtySuffixStr.split(' ')
                    suffixStr = splittedDirtySuffixStr[0]


                else:
                    suffixList = []
                    if suffixNum == 2:
                        # Extract the first suffix and put it in the list
                        dirtySuffixStr = splittedFilter[1]
                        splittedDirtySuffixStr = dirtySuffixStr.split(' ')
                        suffixList.append(splittedDirtySuffixStr[0])

                        # Extract the second suffix and put it in the list
                        dirtySuffixStr = splittedFilter[2]
                        splittedDirtySuffixStr = dirtySuffixStr.split(')')
                        suffixList.append(splittedDirtySuffixStr[0])

                    else:
                        # Extract the first suffix and put it in the list
                        dirtySuffixStr = splittedFilter[1]
                        splittedDirtySuffixStr = dirtySuffixStr.split(' ')
                        suffixList.append(splittedDirtySuffixStr[0])

                        # Extract the last suffix and put it in the list
                        dirtySuffixStr = splittedFilter[suffixNum]
                        splittedDirtySuffixStr = dirtySuffixStr.split(')')
                        suffixList.append(splittedDirtySuffixStr[0])

                        # Extract the rest of the suffixes and put them in the list
                        for i in xrange(2, suffixNum):
                            dirtySuffixStr = splittedFilter[i]
                            splittedDirtySuffixStr = dirtySuffixStr.split(' ')
                            suffixList.append(splittedDirtySuffixStr[0])

                    # Find if the user supplied suffix is valid for the
                    # chosen filetype and set it as the filename suffix
                    isValidSuffix = False
                    userSuffix = '.' + existingSuffix
                    for i in xrange(suffixNum + 1):
                        if userSuffix == suffixList[i]:
                            isValidSuffix = True
                            suffixStr = userSuffix
                            break

                    # If the supplied suffix is not valid replace it
                    # with the default suffix for the chosen filetype
                    if not isValidSuffix:
                        suffixStr = suffixList[0]

            self.outFileName = baseFileName + suffixStr

        if self.inFileName is not None and self.outFileName is not None:
            self.okButton.setDisabled(False)
        self.outputLineEdit.clear()
        self.outputLineEdit.setText(self.outFileName)

    def workerFinished(self, Output_Raster_layer_name):
        # clean up the worker and thread
        self.worker.deleteLater()
        self.thread.quit()
        self.thread.wait()
        self.thread.deleteLater()

        # insert the output raster into QGIS interface
        outputRasterFileInfo = QFileInfo(Output_Raster_layer_name)
        baseName = outputRasterFileInfo.baseName()
        rasterLayer = QgsRasterLayer(Output_Raster_layer_name, baseName)
        if not rasterLayer.isValid():
            print "Layer failed to load"
        QgsMapLayerRegistry.instance().addMapLayer(rasterLayer)

        # remove widget from message bar
        self.close()
        self.thread=None
        self.worker=None

    def workerError(self, e, exception_string):
        QgsMessageLog.logMessage('Worker thread raised an exception:\n'.format(exception_string), level=QgsMessageLog.CRITICAL)

    def startWorker(self):
        if self.thread is not None and self.worker is not None:
            QtGui.QMessageBox.critical(self, "", "the processing is not finished!")
            return
        self.progressBar.setValue(0.0)
        self.outFileName = self.outputLineEdit.text().strip()

        greenBandNo=int(str(self.greenComboBox.currentText()))
        redBandNo=int(str(self.redComboBox.currentText()))
        nirBandNo=int( str(self.nirComboBox.currentText() ) )
        mrBandNo=int( str(self.mRComboBox.currentText() ) )
        if greenBandNo==redBandNo or greenBandNo==nirBandNo or greenBandNo==mrBandNo \
            or redBandNo==nirBandNo or redBandNo==mrBandNo or nirBandNo==mrBandNo:
            QtGui.QMessageBox.critical(self,"", "the input band no is incorrect!")
            return

        # create a new worker instance
        worker = Worker(self.inFileName, self.outFileName,
                        greenNo =greenBandNo,
                        redNo = redBandNo,
                        nirNo= nirBandNo,
                        mrNo = mrBandNo)

        # start the worker in a new thread
        thread = QtCore.QThread(self)
        worker.moveToThread(thread)
        worker.finished.connect(self.workerFinished)
        worker.error.connect(self.workerError)
        worker.progress.connect(self.progressBar.setValue)
        thread.started.connect(worker.run)
        thread.start()
        self.thread = thread
        self.worker = worker